import numpy as np
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l
import sys
sys.path.append('E:\AI\DL')
from Train_ import k_fold as kf
import train_func as tf
def load_data():
    train_data = pd.read_csv('DL\\HousePrice\\data\\train.csv')
    test_data = pd.read_csv('DL\\HousePrice\\data\\test.csv')
    print(train_data.shape)
    all_features = pd.concat((train_data.iloc[:,1:],test_data.iloc[:,1:]))
    print(all_features.shape)
    #连续值归一化,缺项填0
    numeric_features_idx = all_features.dtypes[all_features.dtypes != 'object'].index
    all_features[numeric_features_idx] = all_features[numeric_features_idx]\
        .apply(lambda x:(x-x.mean())/x.std())
    all_features[numeric_features_idx] = all_features[numeric_features_idx].fillna(0)
    #离散值onehot处理
    all_features = pd.get_dummies(all_features,dummy_na=True)
    #nbool无法直接转换为tensor内值
    bool_features_idx = all_features.dtypes[all_features.dtypes == 'bool'].index
    all_features[bool_features_idx] = all_features[bool_features_idx].astype('uint8')
    print(all_features.shape)
    print(all_features.info())

    train_cnt = train_data.shape[0]
    train_features = torch.tensor(all_features[:train_cnt].values,dtype=torch.float32)
    test_features = torch.tensor(all_features[train_cnt:].values,dtype=torch.float32)
    train_labels = torch.tensor(train_data.SalePrice.values.reshape(-1,1),dtype=torch.float32)
    return train_features,test_features,train_labels,test_data

train_features,test_features,train_labels,test_data = load_data()

in_features = train_features.shape[1]
net = nn.Sequential(nn.Linear(in_features,1))

num_epochs = 100
lr = 5
weight_decay = 0
batch_size = 64

def k_fold(k,X_train,y_train):
    train_l_sum,valid_l_sum = 0,0
    for i in range(k):
        data = kf.get_k_fold_data(k,i,X_train,y_train)
        train_ls,valid_ls = tf.train(net,*data,num_epochs,lr,weight_decay,batch_size)
        train_l_sum += train_ls[-1]
        valid_l_sum += valid_ls[-1]
        if i == 0:
            d2l.plot(list(range(num_epochs)),[train_ls,valid_ls],
                     xlabel='epoch',ylabel='mse',xlim=[1,num_epochs],
                     legend=['train','valid'],yscale='log')
        print('fold_num:',i)
        print('train log mse:',train_ls[-1])
        print('valid log mse:',valid_ls[-1])
    d2l.plt.show()
    return train_l_sum / k, valid_l_sum / k
k=5
train_loss,valid_loss = k_fold(k,train_features,train_labels)
print('训练集平均loss:',train_loss)
print('验证集平均loss:',valid_loss)
